'''
Created on Mar 1, 2020
Pytorch Implementation of LightGCN in
Xiangnan He et al. LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation
@author: Jianbai Ye (gusye@mail.ustc.edu.cn)

Design training and test process
'''
import world
import numpy as np
import torch
import utils
import random
import dataloader
from pprint import pprint
from utils import timer
from time import time
from tqdm import tqdm
import model
from parse import args
import multiprocessing
from sklearn.metrics import roc_auc_score


CORES = multiprocessing.cpu_count() // 2


def BPR_train_original(dataset, recommend_model, loss_class, epoch, neg_k=1, w=None):
    Recmodel = recommend_model
    Recmodel.train()
    bpr: utils.BPRLoss = loss_class
    
    with timer(name="Sample"):
        S,Cats = utils.UniformSample_original(dataset)
    users = torch.Tensor(S[:, 0]).long()
    posItems = torch.Tensor(S[:, 1]).long()
    negItems = torch.Tensor(S[:, 2]).long()

    Cats = torch.Tensor(Cats).long()
    Cats = Cats.to(world.device)
    shuffle_indices = np.arange(len(Cats))
    np.random.shuffle(shuffle_indices)
    Cats = Cats[shuffle_indices]

    users = users.to(world.device)
    posItems = posItems.to(world.device)
    negItems = negItems.to(world.device)
    users, posItems, negItems = utils.shuffle(users, posItems, negItems)
    total_batch = len(users) // world.config['bpr_batch_size'] + 1
    vae_batchsize = dataset.vae_datasize//total_batch
    aver_loss = np.array([0.,0.,0.,0.])
    for (batch_i,
         (batch_users,
          batch_pos,
          batch_neg)) in enumerate(utils.minibatch(users,
                                                   posItems,
                                                   negItems,
                                                   batch_size=world.config['bpr_batch_size'])):
        batch_doc = Cats[random.sample(range(len(Cats)), 512)]
        #batch_doc = Cats[vae_batchsize*batch_i:vae_batchsize*batch_i+vae_batchsize]
        cri = bpr.stageOne(batch_users, batch_pos, batch_neg,batch_doc)
        aver_loss += cri
        if world.tensorboard:
            w.add_scalar(f'BPRLoss/BPR', cri[0], epoch * int(len(users) / world.config['bpr_batch_size']) + batch_i)
    aver_loss = aver_loss / total_batch
    world.loss_y.append(aver_loss[0])
    world.pro_logger.logging(f"epoch: {epoch}, loss: {aver_loss[0]:.4f}, bpr: {aver_loss[1]:.4f}, cl: {aver_loss[2]:.4f}, vae: {aver_loss[3]:.4f}")
    print(f"rebult: {sum(bpr.vae.rebult_loss)/len(bpr.vae.rebult_loss):.3f}, kl: {sum(bpr.vae.kl_loss)/len(bpr.vae.rebult_loss):.3f}, reg: {sum(bpr.vae.reg_loss)/len(bpr.vae.rebult_loss):.3f}")
    bpr.vae.reg_loss,bpr.vae.kl_loss,bpr.vae.rebult_loss=[],[],[]
    time_info = timer.dict()
    timer.zero()
    return f"loss{aver_loss[0]:.3f}-{time_info}"
    
    
def test_one_batch(X):
    sorted_items = X[0].numpy()  #batch_size*k
    groundTrue = X[1]  #batch_size*n
    longtailTrue = X[2]
    r = utils.getLabel(groundTrue, sorted_items)
    lr,lp = utils.longtailLabel(longtailTrue,sorted_items)
    pre, recall, ndcg = [], [], []
    tail_pre,tail_recall = [],[]
    for k in world.topks:
        ret = utils.RecallPrecision_ATk(groundTrue, r, k)
        lt_ret = utils.RecallPrecision_LTk(longtailTrue,lr,lp,k)
        pre.append(ret['precision'])
        recall.append(ret['recall'])
        ndcg.append(utils.NDCGatK_r(groundTrue,r,k))
        tail_pre.append(lt_ret['precision'])
        tail_recall.append(lt_ret['recall'])
    return {'recall':np.array(recall),
            'precision':np.array(pre),
            'ndcg':np.array(ndcg),
            'tail_recall':np.array(tail_recall),
            'tail_pre':np.array(tail_pre)}
        
            
def Test(dataset, Recmodel, epoch, w=None, multicore=0):
    u_batch_size = world.config['test_u_batch_size']
    dataset: utils.Loader
    testDict: dict = dataset.testDict
    longtailDict = dataset.longtailDict
    n_tail_user = float(dataset.n_tail_users)
    Recmodel: model.LightGCN
    # eval mode with no dropout
    Recmodel = Recmodel.eval()
    max_K = max(world.topks)
    if multicore == 1:
        pool = multiprocessing.Pool(CORES)
    results = {
        'recall': np.zeros(len(world.topks)),
        'tail_recall': np.zeros(len(world.topks)),
        'precision': np.zeros(len(world.topks)),
        'tail_pre': np.zeros(len(world.topks)),
        'ndcg': np.zeros(len(world.topks))}
    with torch.no_grad():
        users = list(testDict.keys())
        try:
            assert u_batch_size <= len(users) / 10
        except AssertionError:
            print(f"test_u_batch_size is too big for this dataset, try a small one {len(users) // 10}")
        users_list = []
        rating_list = []
        groundTrue_list = []
        longtailTrue_list = []
        # auc_record = []
        # ratings = []
        total_batch = len(users) // u_batch_size + 1
        for batch_users in utils.minibatch(users, batch_size=u_batch_size):
            allPos = dataset.getUserPosItems(batch_users)
            groundTrue = [testDict[u] for u in batch_users]
            longtailTrue = [longtailDict[u] for u in batch_users]
            batch_users_gpu = torch.Tensor(batch_users).long()
            batch_users_gpu = batch_users_gpu.to(world.device)

            rating = Recmodel.getUsersRating(batch_users_gpu)
            #rating = rating.cpu()
            exclude_index = []
            exclude_items = []
            for range_i, items in enumerate(allPos):
                exclude_index.extend([range_i] * len(items))
                exclude_items.extend(items)
            rating[exclude_index, exclude_items] = -(1<<10)
            _, rating_K = torch.topk(rating, k=max_K)

            rating = rating.cpu().numpy()
            # aucs = [ 
            #         utils.AUC(rating[i],
            #                   dataset, 
            #                   test_data) for i, test_data in enumerate(groundTrue)
            #     ]
            # auc_record.extend(aucs)
            del rating
            users_list.append(batch_users)
            rating_list.append(rating_K.cpu())
            groundTrue_list.append(groundTrue)
            longtailTrue_list.append(longtailTrue)
        assert total_batch == len(users_list)
        coverage, tail_coverage, HILD = [], [], []
        for k in world.topks:
            hild = utils.calculate_hmd_average(rating_list, k)
            tail_cov, cov = utils.calculate_long_tail_metrics(rating_list, world.long_tail_doc, k)
            coverage.append(cov)
            HILD.append(hild)
            tail_coverage.append(tail_cov)

        X = zip(rating_list, groundTrue_list,longtailTrue_list)
        if multicore == 1:
            pre_results = pool.map(test_one_batch, X)
        else: #目前是这个
            pre_results = []
            for x in X:
                pre_results.append(test_one_batch(x))
        scale = float(u_batch_size/len(users))
        for result in pre_results:
            results['recall'] += result['recall']
            results['precision'] += result['precision']
            results['ndcg'] += result['ndcg']
            results['tail_pre'] += result['tail_pre']
            results['tail_recall'] += result['tail_recall']
        results['recall'] /= float(len(users))
        results['precision'] /= float(len(users))
        results['ndcg'] /= float(len(users))
        results['tail_pre'] /= float(len(users))
        results['tail_recall'] /= n_tail_user
        results['coverage'] = np.array(coverage)
        results['tail_cov'] = np.array(tail_coverage)
        results['hild'] = np.array(HILD)

        fig_k = 1
        world.recall_y.append(results['recall'][fig_k])
        world.precision.append(results['precision'][fig_k])
        world.ndcg_y.append(results['ndcg'][fig_k])
        world.tail.append(results['tail_pre'][fig_k])
        world.hild.append(HILD[fig_k])
        world.tail_recall.append(results['tail_recall'][fig_k])
        world.tail_coverage.append(results['tail_cov'][fig_k])
        world.coverage.append(results['coverage'][fig_k])
        if world.tensorboard:
            w.add_scalars(f'Test/Recall@{world.topks}',
                          {str(world.topks[i]): results['recall'][i] for i in range(len(world.topks))}, epoch)
            w.add_scalars(f'Test/Precision@{world.topks}',
                          {str(world.topks[i]): results['precision'][i] for i in range(len(world.topks))}, epoch)
            w.add_scalars(f'Test/NDCG@{world.topks}',
                          {str(world.topks[i]): results['ndcg'][i] for i in range(len(world.topks))}, epoch)
        if multicore == 1:
            pool.close()
        print(results)
        return results
